package de.robv.android.xposed;

import static de.robv.android.xposed.XposedHelpers.inputStreamToByteArray;

import android.os.Environment;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;
import java.security.DigestException;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.zip.Adler32;

/**
 * Helper class which can create a very simple .dex file, containing only a class definition with a
 * super class (no methods, fields, ...).
 */
// Dreamland changed: DexCreator is public
public class DexCreator {
  public static File DALVIK_CACHE = new File(Environment.getDataDirectory(), "dalvik-cache");

  /** Returns the default dex file name for the class. */
  public static File getDefaultFile(String childClz) {
    return new File(
        DALVIK_CACHE, "dreamland_" + childClz.substring(childClz.lastIndexOf('.') + 1) + ".dex");
  }

  /**
   * Creates (or returns) the path to a dex file which defines the superclass of {@clz} as extending
   * {@code realSuperClz}, which by itself must extend {@code topClz}.
   */
  public static File ensure(String clz, Class<?> realSuperClz, Class<?> topClz) throws IOException {
    if (!topClz.isAssignableFrom(realSuperClz)) {
      throw new ClassCastException(
          "Cannot initialize " + clz + " because " + realSuperClz + " does not extend " + topClz);
    }

    try {
      return ensure("xposed.dummy." + clz + "SuperClass", realSuperClz);
    } catch (IOException e) {
      throw new IOException("Failed to create a superclass for " + clz, e);
    }
  }

  /** Like {@link #ensure(File, String, String)}, just for the default dex file name. */
  public static File ensure(String childClz, Class<?> superClz) throws IOException {
    return ensure(getDefaultFile(childClz), childClz, superClz.getName());
  }

  /**
   * Makes sure that the given file is a simple dex file containing the given classes. Creates the
   * file if that's not the case.
   */
  public static File ensure(File file, String childClz, String superClz) throws IOException {
    // First check if a valid file exists.
    try {
      byte[] dex = inputStreamToByteArray(new FileInputStream(file));
      if (matches(dex, childClz, superClz)) {
        return file;
      } else {
        file.delete();
      }
    } catch (IOException e) {
      file.delete();
    }

    // If not, create a new dex file.
    byte[] dex = create(childClz, superClz);
    FileOutputStream fos = new FileOutputStream(file);
    fos.write(dex);
    fos.close();
    return file;
  }

  /**
   * Checks whether the Dex file fits to the class names. Assumes that the file has been created
   * with this class.
   */
  public static boolean matches(byte[] dex, String childClz, String superClz) throws IOException {
    boolean childFirst = childClz.compareTo(superClz) < 0;
    byte[] childBytes = stringToBytes("L" + childClz.replace('.', '/') + ";");
    byte[] superBytes = stringToBytes("L" + superClz.replace('.', '/') + ";");

    int pos = 0xa0;
    if (pos + childBytes.length + superBytes.length >= dex.length) {
      return false;
    }

    for (byte b : childFirst ? childBytes : superBytes) {
      if (dex[pos++] != b) {
        return false;
      }
    }

    for (byte b : childFirst ? superBytes : childBytes) {
      if (dex[pos++] != b) {
        return false;
      }
    }

    return true;
  }

  // Dreamland added
  public static ByteBuffer create(String superOf, Class<?> superclass) throws IOException {
    return ByteBuffer.wrap(create("xposed.dummy." + superOf + "SuperClass", superclass.getName()));
  }

  /** Creates the byte array for the dex file. */
  public static byte[] create(String childClz, String superClz) throws IOException {
    boolean childFirst = childClz.compareTo(superClz) < 0;
    byte[] childBytes = stringToBytes("L" + childClz.replace('.', '/') + ";");
    byte[] superBytes = stringToBytes("L" + superClz.replace('.', '/') + ";");
    int stringsSize = childBytes.length + superBytes.length;
    int padding = -stringsSize & 3;
    stringsSize += padding;

    ByteArrayOutputStream out = new ByteArrayOutputStream();

    // header
    out.write("dex\n035\0".getBytes()); // magic
    out.write(new byte[24]); // placeholder for checksum and signature
    writeInt(out, 0xfc + stringsSize); // file size
    writeInt(out, 0x70); // header size
    writeInt(out, 0x12345678); // endian constant
    writeInt(out, 0); // link size
    writeInt(out, 0); // link offset
    writeInt(out, 0xa4 + stringsSize); // map offset
    writeInt(out, 2); // strings count
    writeInt(out, 0x70); // strings offset
    writeInt(out, 2); // types count
    writeInt(out, 0x78); // types offset
    writeInt(out, 0); // prototypes count
    writeInt(out, 0); // prototypes offset
    writeInt(out, 0); // fields count
    writeInt(out, 0); // fields offset
    writeInt(out, 0); // methods count
    writeInt(out, 0); // methods offset
    writeInt(out, 1); // classes count
    writeInt(out, 0x80); // classes offset
    writeInt(out, 0x5c + stringsSize); // data size
    writeInt(out, 0xa0); // data offset

    // string map
    writeInt(out, 0xa0);
    writeInt(out, 0xa0 + (childFirst ? childBytes.length : superBytes.length));

    // types
    writeInt(out, 0); // first type = first string
    writeInt(out, 1); // second type = second string

    // class definitions
    writeInt(out, childFirst ? 0 : 1); // class to define = child type
    writeInt(out, 1); // access flags = public
    writeInt(out, childFirst ? 1 : 0); // super class = super type
    writeInt(out, 0); // no interface
    writeInt(out, -1); // no source file
    writeInt(out, 0); // no annotations
    writeInt(out, 0); // no class data
    writeInt(out, 0); // no static values

    // string data
    out.write(childFirst ? childBytes : superBytes);
    out.write(childFirst ? superBytes : childBytes);
    out.write(new byte[padding]);

    // annotations
    writeInt(out, 0); // no items

    // map
    writeInt(out, 7); // items count
    writeMapItem(out, 0, 1, 0); // header
    writeMapItem(out, 1, 2, 0x70); // strings
    writeMapItem(out, 2, 2, 0x78); // types
    writeMapItem(out, 6, 1, 0x80); // classes
    writeMapItem(out, 0x2002, 2, 0xa0); // string data
    writeMapItem(out, 0x1003, 1, 0xa0 + stringsSize); // annotations
    writeMapItem(out, 0x1000, 1, 0xa4 + stringsSize); // map list

    byte[] buf = out.toByteArray();
    updateSignature(buf);
    updateChecksum(buf);
    return buf;
  }

  private static void updateSignature(byte[] dex) {
    // Update SHA-1 signature
    try {
      MessageDigest md = MessageDigest.getInstance("SHA-1");
      md.update(dex, 32, dex.length - 32);
      md.digest(dex, 12, 20);
    } catch (NoSuchAlgorithmException | DigestException e) {
      throw new RuntimeException(e);
    }
  }

  private static void updateChecksum(byte[] dex) {
    // Update Adler32 checksum
    Adler32 a32 = new Adler32();
    a32.update(dex, 12, dex.length - 12);
    int chksum = (int) a32.getValue();
    dex[8] = (byte) (chksum & 0xff);
    dex[9] = (byte) (chksum >> 8 & 0xff);
    dex[10] = (byte) (chksum >> 16 & 0xff);
    dex[11] = (byte) (chksum >> 24 & 0xff);
  }

  private static void writeUleb128(OutputStream out, int value) throws IOException {
    while (value > 0x7f) {
      out.write((value & 0x7f) | 0x80);
      value >>>= 7;
    }
    out.write(value);
  }

  private static void writeInt(OutputStream out, int value) throws IOException {
    out.write(value);
    out.write(value >> 8);
    out.write(value >> 16);
    out.write(value >> 24);
  }

  private static void writeMapItem(OutputStream out, int type, int count, int offset)
      throws IOException {
    writeInt(out, type);
    writeInt(out, count);
    writeInt(out, offset);
  }

  private static byte[] stringToBytes(String s) throws IOException {
    ByteArrayOutputStream bytes = new ByteArrayOutputStream();
    writeUleb128(bytes, s.length());
    // This isn't MUTF-8, but should be OK.
    bytes.write(s.getBytes(StandardCharsets.UTF_8));
    bytes.write(0);
    return bytes.toByteArray();
  }

  private DexCreator() {}
}
